#include <bits/stdc++.h>
#include "SignAnalyzer.h"
#include "lxmsc.h"

using namespace std;
using namespace llvm;
using namespace SVF;

void SignAnalyzer::init(){
    /// 首先从SVFmodule中获取所有的模块，并为其编号
    this->initBBs();
    /// 从符号表中取出所有的变量
    this->initValues();
}

void SignAnalyzer::proc(){      // the most important function
    int counter = 0;
    while(true){
        bool flag = false;
        /// 对每一个BB依次计算
        for(auto pp: this->bb2id){
            const BB* bb = pp.fi;
            /// 首先根据前驱的BB计算IN值，IN值等于所有前驱的并
            // 判断该BB是否是一个循环结束的BB
            bool flag1 = false;
            string name = (string)bb->getName();
            regex re("while\\.end\\d*");
            regex re1("for\\.end\\d*");
            if (regex_match(name, re) || regex_match(name, re1)) flag1 = true;
            if (flag1) {
                for(auto pred: predecessors(bb)){
                    flag = this->uniteTo(pred, bb, true) || flag;
                }
            }
            else {
                for(auto pred: predecessors(bb)){
                    flag = this->uniteTo(pred, bb, false) || flag;
                }
            }
            
            /// 再根据bb自身的指令计算OUT值
            flag = this->generate(bb) || flag;
        }
        if(this->debug){
            this->report(cout);
            int tmp;
            cin>>tmp;
        }
        counter++;
        if(!flag || counter >= 10) break;
    }
}

bool SignAnalyzer::isInit(const VariableAnsMapTy* block) {
    bool flag = false;
    for (auto p: *block) {
        flag = p.se != INI || flag;
    }
    return flag;
}

void SignAnalyzer::report(ostream &os)const{
    const vector<string> &vec = SignAnalyzer::getSignKindStrings();

    for(auto pp: this->bb2id){
        const BB *bb = pp.fi;
        os<<"===  "<<(string)bb->getName()<<": "<<endl;
        /// 输出符号
        const VariableAnsMapTy &ans = this->bb2ans.find(bb)->se;  // second
        for(auto tt: ans){
            const Value *value = tt.fi;
            if(value->hasName()){
                os<<(string)value->getName()<<": "<<vec[tt.se]<<endl;
            }
        }
    }
}

/// 为所有的BB编号，从1开始
void SignAnalyzer::initBBs(){
    int k = 1;
    for(auto &fun: fanwei(module->llvmFunBegin(), module->llvmFunEnd())){
        auto &s1 = fun->getBasicBlockList();   // get the list of
        for(auto &bb: s1){  // llvm迭代的一般是类本身，bb就是BB类型，不是指针
            this->bb2id.insert(mp(&bb, k++));
            //this->count_ans.insert(mp(&bb, 0));  // 初始化记录BB的数据被更新的次数集合
        }
    }
}

/// 从符号表中获取所有的变量，并且把每个BB都附着一个记录答案的map
void SignAnalyzer::initValues(){
    /// 先做一个去常转型，否则下一句不能通过const检测
    SymbolTableInfo *tabletmp = const_cast<SymbolTableInfo*>(this->table);
    /// m1是一个<const Value*, ID>的map
    const SymbolTableInfo::ValueToIDMapTy &m1 = tabletmp->valSyms();   
    VariableAnsMapTy ret;
    for(auto pp: m1){
        ret.insert(mp(pp.fi, INI));
    } 
    /// 每个BB均保留一个答案记录
    for(auto pp: this->bb2id){
        this->bb2ans.insert(mp(pp.fi, ret));
    }
}

/// 将src的答案合并到tgt上
bool SignAnalyzer::uniteTo(const BB *src, const BB *tgt, bool is_loop){
    if(this->debug){
        cout<<"In funcion: "<<__func__;
        cout<<"("<<(string)src->getName()<<", "<<(string)tgt->getName()<<")"<<endl;   
    }
    bool ret = false;
    /// 取出答案的数据结构
    VariableAnsMapTy &from = this->bb2ans[src];  // 当前BB的前驱
    VariableAnsMapTy &to = this->bb2ans[tgt];  // 当前的BB
    /// 对每一个变量的答案
    vartovar var_map;   // 临时的寄存器变量 到 原始变量  的映射
    CmpInst::Predicate sign;
    llvm::Value* op1, *op2;
    int count1 = 0;
    int count2 = 0;
    if (is_loop) {
        for(auto &inst: beend(src)){  // 迭代出来的是指令本身，所以要取地址
            if (inst.getOpcode() == Instruction::Load) {
                if (checkForBr(src, inst.getOperand(0))) {
                    var_map.insert(mp((const llvm::Value*)inst.getOperandUse(0).getUser(), inst.getOperand(0)));   // %0 -> &x
                }
            }
            if (inst.getOpcode() == Instruction::ICmp) {
                sign = ((const llvm::CmpInst &)inst).getPredicate();
                if (var_map.find(inst.getOperand(0)) == var_map.end()) op1 = inst.getOperand(0);
                else {op1 = var_map.find(inst.getOperand(0))->se; count1++;}
                if (var_map.find(inst.getOperand(1)) == var_map.end()) op2 = inst.getOperand(1);
                else {op2 = var_map.find(inst.getOperand(1))->se; count2++;}
            }
        }
    }
    for(auto &pp: to){  // 使用引用，要修改值
        SignK tmp;
        if (is_loop) {
            if (count1 >= 1) {
                if (((string)op1->getName()) == ((string)pp.fi->getName())) {
                    tmp = doLoop(sign, op2,from, false);
                }
                else {
                    //cout << "8" << endl;
                    tmp = SignAnalyzer::unite(from[pp.fi], pp.se);
                }
            }
            else if (count2 >= 1) {
                if (((string)op2->getName()) == ((string)pp.fi->getName())) {
                    tmp = doLoop(sign, op1,from, true);
                }
                else {
                    tmp = SignAnalyzer::unite(from[pp.fi], pp.se);
                }
            }
            else {
                //cout << "3" << endl;
                tmp = SignAnalyzer::unite(from[pp.fi], pp.se);
            }
        } else {
            //cout << "9" << endl;
            tmp = SignAnalyzer::unite(from[pp.fi], pp.se);
        }
        //cout << getSignKindString(tmp) << endl;
        //cout << getSignKindString(pp.se) << endl;
        // if (is_loop) {
        //     cout << (string)pp.fi->getName() << endl;
        //     cout << (string)op1->getName() << endl;
        // }
        if(tmp != pp.se){
            if(this->debug){
                cout<<"value update occured: unite("<<from[pp.fi]<<", "<<pp.se<<"), ";
                cout<<"("<<pp.fi<<", "<<(string)pp.fi->getName()<<")";
                cout<<", oldvalue = "<<pp.se<<", newvalue = "<<tmp<<endl;
                
            }            
            pp.se = tmp;
            //cout << getSignKindString(pp.se) << endl;
            if (count1 >= 1 || count2 >= 1) {
                ret = false;
            }else {
                ret = true;
            }
            if(this->debug){
                cout<<"after: "<<to[pp.fi]<<endl;
            }
        }
    }
    return ret;
}

/// 根据BB的指令计算答案
bool SignAnalyzer::generate(const BB *bb){
    if(this->debug){
        cout<<"In funcion: "<<__func__;
        cout<<"("<<(string)bb->getName()<<")"<<endl;
    }
    bool flag = false;
    VariableAnsMapTy &ans = this->bb2ans[bb];  
    for(auto &inst: beend(bb)){  // 迭代出来的是指令本身，所以要取地址
        flag = this->generate(&inst, ans, bb) || flag;
    }
    if(this->debug){
        cout<<"return: "<<flag<<endl;
    }
    return flag;
}

/**
 * 根据指令，更新ans中的内容
 * 指令的枚举即Instruction:Store等的定义在
 * llvm/IR/Instruction.def
 * 文件中
*/
bool SignAnalyzer::generate(const Instruction *inst, VariableAnsMapTy &ans, const BB* bb){
    bool flag = false;
    switch(inst->getOpcode()){
        case Instruction::Store:{ // store 值, 目标，值有可能是常数，也有可能是变量
            const Value *op0 = inst->getOperand(0);
            const Value *op1 = inst->getOperand(1);
            flag = this->genStore(op0, op1, ans) || flag;
        }break;
        case Instruction::Load:{  // user = load op0
            const Value *op0 = inst->getOperand(0);
            const User *user = inst->getOperandUse(0).getUser();  
            flag = this->genLoad(user, op0, ans) || flag;
        }break;
        case Instruction::Add:{  // user = add op0, op1
            const Value *op0 = inst->getOperand(0);
            const Value *op1 = inst->getOperand(1);
            const User *user = inst->getOperandUse(0).getUser(); 
            flag = this->genAdd(user, op0, op1, ans, bb) || flag;
        }break;
        case Instruction::SDiv:{  // 如果是除法指令
            const Value *op0 = inst->getOperand(0);
            const Value *op1 = inst->getOperand(1);
            const User *user = inst->getOperandUse(0).getUser(); 
            flag = this->genSdiv(user, op0, op1, ans, bb) || flag;
        }break;
        case Instruction::Mul:{  // 如果是乘法指令
            const Value *op0 = inst->getOperand(0);
            const Value *op1 = inst->getOperand(1);
            const User *user = inst->getOperandUse(0).getUser(); 
            flag = this->genMul(user, op0, op1, ans, bb) || flag;
        }break;
        case Instruction::Sub:{  // 如果是减法指令
            const Value *op0 = inst->getOperand(0);
            const Value *op1 = inst->getOperand(1);
            const User *user = inst->getOperandUse(0).getUser(); 
            flag = this->genSub(user, op0, op1, ans, bb) || flag;
        }break;
    }
    return flag;
}

/// 处理store指令，store op0, op1， op0有可能是常数，op1必然是地址
bool SignAnalyzer::genStore(const Value *op0, const Value *op1, VariableAnsMapTy &ans){
    SignK ret = this->getSign(op0, ans);
    SignK &target = ans[op1];
    return ret != target ? (target=ret, true) : false;
}

/// user = load op0，op0是地址
bool SignAnalyzer::genLoad(const User *user, const Value *op0, VariableAnsMapTy &ans){
    SignK ret = ans[op0];
    SignK &target = ans[user];
    return ret != target ? (target=ret, true) : false;
}

/// user = add op0, op1
bool SignAnalyzer::genAdd(const User *user, const Value *op0, const Value *op1, VariableAnsMapTy &ans, const BB* bb){
    SignK a = this->getSign(op0, ans);
    SignK b = this->getSign(op1, ans);
    SignK ret = INI;
    switch(a){
        case POS:ret = (POS==b||ZER==b) ? POS : UNK;break;
        case ZER:ret = b;break;
        case NEG:ret = (NEG==b||ZER==b) ? NEG : UNK;break;
        case UNK:ret = UNK;break;
    }
    SignK &target = ans[user];
    return ret != target ? (target=ret, true) : false;     // judge the variable user 
}

/// user = sdiv op0, op1
bool SignAnalyzer::genSdiv(const User *user, const Value *op0, const Value *op1, VariableAnsMapTy &ans, const BB* bb){
    SignK a = this->getSign(op0, ans);
    SignK b = this->getSign(op1, ans);
    SignK ret = INI;
    switch(a){
        case POS:ret = (POS==b||NEG==b) ? b : UNK;break;
        case ZER:ret = ZER;break;
        case NEG:ret = NEG==b ? POS : (POS==b ? NEG : UNK);break;
        case UNK:ret = UNK;break;
    }
    SignK &target = ans[user];
    return ret != target ? (target=ret, true) : false;
}

/// user = sub op0, op1
bool SignAnalyzer::genSub(const User *user, const Value *op0, const Value *op1, VariableAnsMapTy &ans, const BB* bb) {
    SignK a = this->getSign(op0, ans);
    SignK b = this->getSign(op1, ans);
    SignK ret = INI;
    switch(a){
        case POS:ret = (ZER==b||NEG==b) ? POS : UNK;break;
        case ZER:ret = (b == ZER) ? ZER : (b == POS ? NEG : (b == NEG ? POS : UNK));break;
        case NEG:ret = (b == POS || b == ZER) ? NEG : UNK;break;
        case UNK:ret = UNK;break;
    }
    SignK &target = ans[user];
    return ret != target ? (target=ret, true) : false;
}

/// user = mul op0, op1
bool SignAnalyzer::genMul(const User *user, const Value *op0, const Value *op1, VariableAnsMapTy &ans, const BB* bb){
    SignK a = this->getSign(op0, ans);
    SignK b = this->getSign(op1, ans);
    SignK ret = INI;
    switch(a){
        case POS:ret = (POS==b||NEG==b || ZER == b) ? b : UNK;break;
        case ZER:ret = ZER;break;
        case NEG:ret = (NEG == b) ? POS : ((POS == b) ? NEG : ((ZER == b) ? ZER : UNK));break;
        case UNK:ret = UNK;break;
    }
    SignK &target = ans[user];
    return ret != target ? (target=ret, true) : false;
}

const vector<string>&  SignAnalyzer::getSignKindStrings(){
    static const int n = 5;
    static const string a[n] = {
        "INI", "POS", "NEG", "ZER", "UNK"
    };
    static const vector<string> vec(a, a+n);
    return vec;
}

string SignAnalyzer::getSignKindString(SignK k){
    return SignAnalyzer::getSignKindStrings()[k];
}

// 检查变量是否没有变化
bool SignAnalyzer::checkForBr(const BB* bb, const llvm::Value * var) {
    bool flag = false;
    VariableAnsMapTy &ans = bb2ans.find(bb)->se;
    for (auto pre: predecessors(bb)) {
        for (auto tt: bb2ans.find(pre)->se) {
            if ((string)tt.fi->getName() == (string)var->getName() && tt.se != getSign(var, ans) && tt.se != INI && getSign(var, ans) != INI) {
                flag = true;  // 表示这个这个变量是条件判定值
            }
        }
    }
    return flag;
}

SignAnalyzer::SignK SignAnalyzer::doLoop(llvm::CmpInst::Predicate cond, const llvm::Value* op2, VariableAnsMapTy &ans, bool inv) {

    if (inv) {
        switch(cond) {
            case llvm::CmpInst::ICMP_SGT:
                cond = llvm::CmpInst::ICMP_SLT;
                break;
            case llvm::CmpInst::ICMP_SLT:
                cond = llvm::CmpInst::ICMP_SGT;
                break;
            case llvm::CmpInst::ICMP_SGE:
                cond = llvm::CmpInst::ICMP_SLE;
                break;
            case llvm::CmpInst::ICMP_SLE:
                cond = llvm::CmpInst::ICMP_SGE;
                break;
            default:
                break;
        }
    }

    switch(cond) {
        case llvm::CmpInst::ICMP_EQ:{
            return UNK;
        }break;
        case llvm::CmpInst::ICMP_NE:{
            return getSign(op2, ans);
        }break;
        case llvm::CmpInst::ICMP_SGT:{
            if (getSign(op2, ans) == NEG) {return NEG;}
            else {return UNK;}
        }break;
        case llvm::CmpInst::ICMP_SLT:{
            if (getSign(op2, ans) == POS) return POS;
            else return UNK;
        }break;
        case llvm::CmpInst::ICMP_SGE:{
            if (getSign(op2, ans) == NEG || getSign(op2, ans) == ZER) return NEG;
            else return UNK;
        }break;
        case llvm::CmpInst::ICMP_SLE:{
            if (getSign(op2, ans) == POS || getSign(op2, ans) == ZER) return POS;
            else return UNK;
        }break;
        default:
            return UNK;
    }
}